package com.zjj.lbw.ai.vector;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.redis.RedisEmbeddingStore;

import java.util.List;

public class VectorSearch {

    public static void main(String[] args) {

        OpenAiEmbeddingModel embeddingModel = OpenAiEmbeddingModel.builder()
                .baseUrl("http://langchain4j.dev/demo/openai/v1")
                .apiKey("demo")
                .modelName("text-embedding-3-small")
                .build();

        RedisEmbeddingStore embeddingStore = RedisEmbeddingStore.builder()
                .host("127.0.0.1")
                .port(6379)
                .dimension(1536)
                .build();

        // 生成向量
//        TextSegment segment1 = TextSegment.from("我是子坚君");
//        Response<Embedding> embed1 = embeddingModel.embed(segment1);
//
//        // 存储向量
//        embeddingStore.add(embed1.content(), segment1);

        // 生成向量
        Response<Embedding> embed2 = embeddingModel.embed("我的名字叫子坚君");
//        Response<Embedding> embed2 = embeddingModel.embed("今天天气很好");

        List<EmbeddingMatch<TextSegment>> result = embeddingStore.findRelevant(embed2.content(), 4);
        for (EmbeddingMatch<TextSegment> embeddingMatch : result) {
            System.out.println(embeddingMatch.embedded().text());
            System.out.println(embeddingMatch.score());
        }
    }
}
